// Copyright 2019 DeepMind Technologies Limited.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef REVERB_CC_TIME_MONITOR_H_
#define REVERB_CC_TIME_MONITOR_H_

#include "absl/time/time.h"
#include "reverb/cc/platform/hash_map.h"

namespace deepmind {
namespace reverb {
namespace internal {

// Keeps track of how much time is spent in each state. Not thread safe.
template <typename StateEnum> class StateStatistics {
 public:
  // Enters a given state.
  void Enter(StateEnum new_state) {
    auto now = absl::Now();
    if (last_state_change_time_ != absl::InfiniteFuture()) {
      time_per_state_[active_state_] += now - last_state_change_time_;
    }
    last_state_change_time_ = now;
    active_state_ = new_state;
  }

  // Provides aggregated time spent in a given state.
  absl::Duration GetTotalTimeIn(StateEnum state) const {
    absl::Duration duration = absl::ZeroDuration();
    auto it = time_per_state_.find(state);
    if (it != time_per_state_.end()) {
      duration += it->second;
    }
    if (active_state_ == state &&
        last_state_change_time_ != absl::InfiniteFuture()) {
      duration += absl::Now() - last_state_change_time_;
    }
    return duration;
  }

  // Returns the current state.
  StateEnum CurrentState() const {
    return active_state_;
  }

 private:
  StateEnum active_state_ = StateEnum(0);
  absl::Time last_state_change_time_ = absl::InfiniteFuture();
  internal::flat_hash_map<StateEnum, absl::Duration> time_per_state_;
};

}  // namespace internal
}  // namespace reverb
}  // namespace deepmind

#endif  // REVERB_CC_TIME_MONITOR_H_
